﻿using FastApi.Attributes;

namespace FastApi
{
    public class FastRateLimit
    {
        private readonly Dictionary<string, SemaphoreSlim> _dict = [];

        private readonly object _locker = new();

        private SemaphoreSlim Get(string key, int maxCount)
        {
            var ok = _dict.TryGetValue(key, out var val);
            if (!ok)
            {
                lock (_locker)
                {
                    if (!_dict.ContainsKey(key))
                    {
                        _dict.Add(key, new SemaphoreSlim(maxCount));
                    }
                    val = _dict[key];
                }
            }
            return val;
        }

        /// <summary>
        /// 等待一个
        /// </summary>
        /// <param name="ctx"></param>
        /// <param name="key"></param>
        /// <param name="maxCount">允许最大并发数</param>
        /// <returns></returns>
        public bool WaitOne(FastContext ctx, string key, int maxCount, string errMsg = "FastRateLimit", int errCode = -1)
        {
            var val = Get(key, maxCount);
            var ok = val.Wait(0);
            if (ok)
            {
                ctx.UserData.Add(key, null);
            }
            else
            {
                ctx.Ex = new FastException(errMsg, errCode);
            }
            return ok;
        }

        /// <summary>
        /// 释放
        /// </summary>
        /// <param name="ctx"></param>
        /// <param name="key"></param>
        /// <returns></returns>
        public bool Release(FastContext ctx, string key)
        {
            if (ctx.UserData.ContainsKey(key))
            {
                _dict[key].Release();
                return true;
            }
            return false;
        }

    }
}
